// "Copyright [2021] <Copyright Shengkai Lin>"

#include <string>
#include <vector>
#include "network_routing.h"

namespace ai_sim {
Define_Module(NetworkRouting);

void NetworkRouting::initialize() {
    cModule * topo_module = getModuleByPath("^.^.topo_manager");
    topos = omnetpp::check_and_cast<TopoManager *> (topo_module);
    needStaticLeaf = par("needStaticLeaf");
    clustertype = ClusterType(getSystemModule()->par("clustertype").intValue());
    EV_INFO << "INIT NetworkRouting FINISHED" << endl;
}

bool NetworkRouting::selectByName(cModule *mod, void* name_vec) {
    string fullpath = mod->getFullPath();
    vector<string> *name_vec_ = (vector<string>*) name_vec;
    for (string name_ : *name_vec_) {
        auto pos = fullpath.find(name_);
        if (pos != string::npos) {
            // Should make sure the submodules is not selected
            if (fullpath.substr(pos).find(".") == string::npos) return true;
        }
    }
    return false;
}

void NetworkRouting::gen_route(
    taskboard& tmp_taskboard, RouteAlgorithm* algo) {
    Enter_Method_Silent();
////////////////////// Generate and allocate routes/////////////////////////
// Find the shortest path of the comm_pair one by one
// The time complexity of algorithm is pretty HIGH.
    algo->generateRoutes(tmp_taskboard, topos);
    EV_INFO << "Finished gen routes" << std::endl;
}

void NetworkRouting::gen_txt_route(
    map<int, RouteVec>& route_in_switch, DetailFlow& flow) {
    Topology* dataplane_topo = topos->getWorkersTopo();
    int inport = -1;
    for (auto comp_it = flow.flow_path.begin();
        comp_it != flow.flow_path.end(); comp_it ++) {
        auto next_it = comp_it + 1;
        // End of the path
        if (next_it == flow.flow_path.end()) break;
        // get src node ptr
        Topology level_topo;
        std::vector<std::string> ned_vec;
        ned_vec.push_back(comp_it->modname);
        level_topo.extractFromNetwork(selectByName, (void *)&ned_vec);

        cModule* src_mod =  level_topo.getNode(comp_it->index)->getModule();
        EV << "***************************" << endl;
        EV << "src name " << src_mod->getFullPath() << endl;
        comp_it->modid = src_mod->getId();
        // get dest node ptr
        ned_vec.clear();
        ned_vec.push_back(next_it->modname);
        level_topo.clear();
        level_topo.extractFromNetwork(selectByName, (void *)&ned_vec);

        cModule* dest_mod =  level_topo.getNode(next_it->index)->getModule();
        EV << "dest mod " << dest_mod->getFullPath() << endl;
        // Get both level topo
        auto dp_topo = topos->getDataplaneTopo();
        EV << level_topo.getNumNodes() << endl;
        auto* src_node = dp_topo->getNodeFor(src_mod);
        auto* dest_node = dp_topo->getNodeFor(dest_mod);
        vector<int> avail_path;
        EV << "src_node->getNumOutLinks()" << src_node->getNumOutLinks() << endl;
        EV << "connect to:" << endl;
        for (int i = 0; i < src_node->getNumOutLinks(); i++) {
            EV << src_node->getLinkOut(i)->getRemoteNode()->getModule()->getFullPath() << endl;
            if (src_node->getLinkOut(i)->getRemoteNode() == dest_node)
                avail_path.push_back(i);
        }
        // Get route info
        if (comp_it->outpathidx >= avail_path.size()) {
            EV << "avail_path.size()" << avail_path.size() << endl;
            EV << "comp_it->outpathidx" << comp_it->outpathidx << endl;
            error("Link cnt not enough");
        } else {
            auto* choose_link = src_node->getLinkOut(
                avail_path[comp_it->outpathidx]);
            int outport = choose_link->getLocalGateId();

            RouteInfo now_route = {
                .inPort = inport,
                .outPort = outport,
                .srcIdx = flow.flow_path.front().index,
                .destIdx = flow.flow_path.back().index
            };
            route_in_switch[src_node->getModuleId()].push_back(now_route);
            inport = choose_link->getRemoteGateId();
        }
    }
    EV << "Gen path finished" << endl;
}
void NetworkRouting::allocRoutes(const map<int, RouteVec>& routes_map) {
    Enter_Method_Silent();
    for (auto tmp_routes : routes_map) {
        if (topos->getControlSdnTopo()->getNodeFor(
            getSimulation()->getModule(tmp_routes.first)) == nullptr) {
            // If the node is not the switch (for example nics)
            continue;
        }
            visitedLeaf[tmp_routes.first] = true;
            RouteVecMsg *ra_msg = new RouteVecMsg;
            ra_msg->setKind(ROUTE_ALLOC_MSG);
            ra_msg->setNew_route_table(tmp_routes.second);
            ra_msg->setTargetId(tmp_routes.first);
            // Get the output port to the machine the Gpu in.
            // If cluster type is ElectricalCluster, controller will send route to
            // adapter instead of oxc or switch directly

            // RouteVec temp_routes = tmp_routes.second;
            // std::ofstream ofile;
            // for(auto temp_route : temp_routes) {
            //     ofile.open("./record_route.txt",ios::app);
            //     ofile << tmp_routes.first << "-" << temp_route.srcIdx << "-" << temp_route.destIdx << "-" << temp_route.inPort << "-" << temp_route.outPort <<endl;
            // }
            int outport = topos->getControlSdnTopo()->getNodeFor(
                getSimulation()->getModule(tmp_routes.first))->
                getinPathfromDest(0)->getRemoteGateId();

            EV << "The remote gate is " << topos->getControlSdnTopo()->
                getNodeFor(getSimulation()->getModule(tmp_routes.first))->
                getinPathfromDest(0)->getRemoteGate()->getFullPath()<<
                std::endl;

            send(ra_msg, outport);
            EV << "Alloc a route to " << getSimulation()->
                getModule(tmp_routes.first)->getFullPath() << std::endl;

    }
    EV << "Finished alloc routes on start" << endl;
}

void NetworkRouting::alloctxtRoutes(const map<int, RouteVec>& routes_map) {
    Enter_Method_Silent();
    for (auto tmp_routes : routes_map) {
        if (topos->getControlSdnTopo()->getNodeFor(
            getSimulation()->getModule(tmp_routes.first)) == nullptr) {
            // If the node is not the switch (for example nics)
            continue;
        }
        RouteVecMsg *ra_msg = new RouteVecMsg;
        ra_msg->setKind(ROUTE_ALLOC_MSG);
        ra_msg->setNew_route_table(tmp_routes.second);
        ra_msg->setTargetId(tmp_routes.first);
        // Get the output port to the machine the Gpu in.
        // If cluster type is ElectricalCluster, controller will send route to
        // adapter instead of oxc or switch directly
        int outport = topos->getControlSdnTopo()->getNodeFor(
            getSimulation()->getModule(tmp_routes.first))->
            getinPathfromDest(0)->getRemoteGateId();

        EV << "The remote gate is " << topos->getControlSdnTopo()->
            getNodeFor(getSimulation()->getModule(tmp_routes.first))->
            getinPathfromDest(0)->getRemoteGate()->getFullPath()<<
            std::endl;

        send(ra_msg, outport);
        EV << "Alloc a route to " << getSimulation()->
            getModule(tmp_routes.first)->getFullPath() << std::endl;
    }
    EV << "Finished alloc routes on start" << endl;
}

void NetworkRouting::allocRoutes(FlowMsg* f_msg) {
    Enter_Method_Silent();
    const map<int, RouteVec>& routes_map = f_msg->getThe_flow().dependRoute;
    EV << "Size of depend routes" << routes_map.size() << endl;
    for (auto tmp_routes : routes_map) {
        auto* switch_node = topos->getControlSdnTopo()->getNodeFor(
            getSimulation()->getModule(tmp_routes.first));
        if (switch_node == nullptr) {
            // If the node is not the switch (for example nics)
            continue;
        }
        int outport = topos->getControlSdnTopo()->getNodeFor(
            getSimulation()->getModule(tmp_routes.first))->
            getinPathfromDest(0)->getRemoteGateId();
        const char * switch_type_char =
            switch_node->getModule()->getNedTypeName();
        string switch_type = string(switch_type_char);
        // EV << switch_type << endl;
        if (switch_type == "inet.node.inet.Router") {
            auto targetRouter = dynamic_cast<inet::Router *>(getSimulation()->getModule(tmp_routes.first));
            RouteVecMsg *ra_msg = new RouteVecMsg;
            ra_msg->setKind(ROUTE_ALLOC_MSG);
            ra_msg->setNew_route_table(tmp_routes.second);
            ra_msg->setTargetId(tmp_routes.first);
            ra_msg->setIndex(targetRouter->getIndex());
            ra_msg->setThe_flow(f_msg->getThe_flow());
            record_flow_count[std::make_pair(
                (const int&)f_msg->getThe_flow().srcIdx,
                (const int&)f_msg->getThe_flow().destIdx)]++;
            send(ra_msg, outport);
            SwitchAckMsg *oxc_ack_msg = new SwitchAckMsg;
            oxc_ack_msg->setThe_flow(f_msg->getThe_flow());
            oxc_ack_msg->setKind(SWITCH_ACK_MSG);
            scheduleAt(omnetpp::simTime(), oxc_ack_msg);
        } else if (switch_type == "cluster.optical_devices.OpticalDevice") {
            // check whether the path to be added is valid,
            // and get those path which can be accessed
            std::vector<Relay> path;
            for (auto ro : tmp_routes.second) {
                Relay new_relay = Relay(tmp_routes.first, ro.inPort, ro.outPort);
                path.push_back(new_relay);
            }
            Flow the_flow = f_msg->getThe_flow();
            optical_path_manager_.AddPath(
                the_flow.srcIdx, the_flow.destIdx, path);
            ChangeSet change_set =
                optical_path_manager_.GetChangeSet(tmp_routes.first);
            record_flow_count[std::make_pair(
                the_flow.srcIdx, the_flow.destIdx)]++;
            OXCConfMsg *oxc_msg = new OXCConfMsg;
            oxc_msg->setChange_set(change_set);
            oxc_msg->setThe_flow(f_msg->getThe_flow());
            oxc_msg->setKind(OXC_CONF_MSG);
            send(oxc_msg, outport);
        } else if (switch_type == "src.IBSwitch") {
            RouteVecMsg *ra_msg = new RouteVecMsg;
            ra_msg->setKind(ROUTE_ALLOC_MSG);
            ra_msg->setNew_route_table(tmp_routes.second);
            ra_msg->setTargetId(tmp_routes.first);
            send(ra_msg, outport);
            EV << "The remote gate is " << topos->getControlSdnTopo()->
                getNodeFor(getSimulation()->getModule(tmp_routes.first))->
                getinPathfromDest(0)->getRemoteGate()->getFullPath()<<
                std::endl;
            EV << "Alloc a route to " << getSimulation()->
                getModule(tmp_routes.first)->getFullPath() << std::endl;
        } else if (switch_type == "cluster.optical_devices.OCS") {
            // check whether the path to be added is valid,
            // and get those path which can be accessed
            std::vector<Relay> path;
            for (auto ro : tmp_routes.second) {
                Relay new_relay = Relay(tmp_routes.first, ro.inPort, ro.outPort);
                path.push_back(new_relay);
            }
            Flow the_flow = f_msg->getThe_flow();
            optical_path_manager_.AddPath(
                the_flow.srcIdx, the_flow.destIdx, path);
            ChangeSet change_set =
                optical_path_manager_.GetChangeSet(tmp_routes.first);
            record_flow_count[std::make_pair(
                the_flow.srcIdx, the_flow.destIdx)]++;
            OXCConfMsg *oxc_msg = new OXCConfMsg;
            oxc_msg->setChange_set(change_set);
            oxc_msg->setThe_flow(f_msg->getThe_flow());
            oxc_msg->setKind(OXC_CONF_MSG);
            send(oxc_msg, outport);
        } else {
            throw omnetpp::cRuntimeError("Invalid switch type!");
        }
    }
    EV << "Finished alloc routes from msg" << endl;
}

void NetworkRouting::handleMessage(omnetpp::cMessage *msg) {
    Enter_Method_Silent();
    SwitchAckMsg * oxc_ack_msg = dynamic_cast<SwitchAckMsg *> (msg);
    if (msg->getKind() == SYNC_MSG) {
        error("DONT SUPPORT SYNC MSG YET");
        // FlowMsg *sy_msg = dynamic_cast<FlowMsg *>(msg);
        // send(msg, "machine_port$o", sy_msg->getDestIdx()%numGpusperMachine);
    } else if (msg->getKind() == SWITCH_ACK_MSG || oxc_ack_msg != nullptr) {
        auto flow_key = std::make_pair(
            oxc_ack_msg->getThe_flow().srcIdx,
            oxc_ack_msg->getThe_flow().destIdx);
        if (record_flow_count[flow_key] != 0) {
            record_flow_count[flow_key]--;
            if (record_flow_count[flow_key] == 0) {
                msg->setKind(SWITCH_ACK_MSG);
                send(msg, "dispatcher_port");
            } else {
                delete oxc_ack_msg;
            }
        } else {
            delete oxc_ack_msg;
        }
    } else {
        delete msg;
        error("ERROR MSG ARRIVED ON CONTROLLER");
    }
}
}
